#include <string.h>

#define _WIN32_WINNT 0x0500
#include <Windows.h>
#include <Tlhelp32.h>
#include <stdlib.h>

#include "include/log.h"
#include "include/botdata.h"
#include "include/envhook.h"
#include "include/pointers.h"

void *EnvReturnAddress;
void *EnvHookBridgeAddress;
static struct BotData *botDataHook;

void EnvHookBody()
{
    /* Wait for our mutex */
    if(WaitForSingleObject(botDataHook->environment.mutex, INFINITE) != WAIT_OBJECT_0)
    {
        MessageBox(NULL, "WaitForSingleObject failed", NULL, MB_OK | MB_ICONERROR);
        exit(1);
    }

    /* Check the current array size and allocate new if needed */
    if(*ENVIRONMENT_MAX_PTR > botDataHook->environment.limit)
    {
        free(botDataHook->environment.object);
        botDataHook->environment.object = malloc(*ENVIRONMENT_MAX_PTR * sizeof(struct RomObject *));
        if(botDataHook->environment.object == NULL)
        {
            MessageBox(NULL, "Out of memory", NULL, MB_OK | MB_ICONERROR);
            exit(1);
        }
    }

    /* Copy over the contents */
    botDataHook->environment.limit = *ENVIRONMENT_MAX_PTR;
    botDataHook->environment.size = *ENVIRONMENT_SIZE_PTR;
    memcpy(botDataHook->environment.object, *ENVIRONMENT_PTR, *ENVIRONMENT_SIZE_PTR * sizeof(struct RomObject *));
    
    if(ReleaseMutex(botDataHook->environment.mutex) == 0)
    {
        MessageBox(NULL, "ReleaseMutex failed", NULL, MB_OK | MB_ICONERROR);
        exit(1);
    }

    return;
}

int EnvHookInstall(struct BotData *botData)
{
    LogWrite(botData, "Attempting to install environment hooks...");

    /* Set all global variables */
    EnvReturnAddress = (void*)0x0061B95C;
    EnvHookBridgeAddress = &EnvHookBridge;
    botDataHook = botData;
    
    /* Enumerate all threads */
    HANDLE threadSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPTHREAD, 0);
    if(threadSnapshot == INVALID_HANDLE_VALUE)
    {
        LogWrite(botData, "Error: CreateToolHelp32Snapshot failed with %lx", GetLastError());
        return 1;
    }
    
    THREADENTRY32   te32;
    te32.dwSize = sizeof(THREADENTRY32);
    
    if(Thread32First(threadSnapshot, &te32) == 0)
    {
        LogWrite(botData, "Error: Thread32First failed with %lX", GetLastError());
        CloseHandle(threadSnapshot);
        return 1;
    }
    
    unsigned long ourPid = GetCurrentProcessId();
    unsigned long ourTid = GetCurrentThreadId();
    do
    {
        /* Suspend all threads that aren't this thread */
        if(te32.th32OwnerProcessID == ourPid && te32.th32ThreadID != ourTid)
        {
            HANDLE threadHandle;
            threadHandle = OpenThread(THREAD_SUSPEND_RESUME, 0, te32.th32ThreadID);
            if(threadHandle == INVALID_HANDLE_VALUE)
            {
                LogWrite(botData, "Error: OpenThread failed with %lX", GetLastError());
                CloseHandle(threadSnapshot);
                return 1;
            }
            if(SuspendThread(threadHandle) == (unsigned long)-1) /* Nice job M$, returning -1 unsigned long... */
            {
                LogWrite(botData, "Error: SuspendThread failed with %lX", GetLastError());
                CloseHandle(threadHandle);
                CloseHandle(threadSnapshot);
                return 1;
            }
            CloseHandle(threadHandle);
        }
    }while(Thread32Next(threadSnapshot, &te32));

    /* 89 6E 0C 8B 46 0C*/
    /* Patch in the jump */
    unsigned long oldProtect;
    VirtualProtect((void *)0x0061B956, 6, PAGE_EXECUTE_READWRITE, &oldProtect);
    memcpy((void *)0x0061B956, "\xFF\x25", 2);
    void **patchMe = (void **)0x0061B958;
    *patchMe = (void *)&EnvHookBridgeAddress;
    VirtualProtect((void *)0x0061B956, 6, oldProtect, &oldProtect);

    /* Resume all threads*/
    if(Thread32First(threadSnapshot, &te32) == 0)
    {
        LogWrite(botData, "Error: Thread32First failed with %lX", GetLastError());
        CloseHandle(threadSnapshot);
        return 1;
    }
    
    do
    {
        /* Resume all threads that aren't this thread */
        if(te32.th32OwnerProcessID == ourPid && te32.th32ThreadID != ourTid)
        {
            HANDLE threadHandle = OpenThread(THREAD_SUSPEND_RESUME, 0, te32.th32ThreadID);
            if(threadHandle == INVALID_HANDLE_VALUE)
            {
                LogWrite(botData, "Error: OpenThread failed with %lX", GetLastError());
                CloseHandle(threadSnapshot);
                return 1;
            }
            if(ResumeThread(threadHandle) == (unsigned long)-1) /* Nice job M$, returning -1 unsigned long... */
            {
                LogWrite(botData, "Error: SuspendThread failed with %lX", GetLastError());
                CloseHandle(threadHandle);
                CloseHandle(threadSnapshot);
                return 1;
            }
            CloseHandle(threadHandle);
        }
    }while(Thread32Next(threadSnapshot, &te32));

    CloseHandle(threadSnapshot);
    return 0;
}
